#!/usr/bin/env python
# -*- coding: UTF-8 -*-
"""
@Project ：OA-DG-main 
@File    ：faster_rcnn_r50_fpn_1x_thyroid.py
@IDE     ：PyCharm 
@Author  ：cao xu
@Date    ：2025/9/16 下午3:17 
"""
# configs/OA-DG/thyroid/faster_rcnn_r50_fpn_1x_thyroid.py
_base_ = [
    '/home/ai999/project/OA-DG/configs/_base_/models/faster_rcnn_r50_fpn.py',
    '/home/ai999/project/OA-DG/configs/_base_/default_runtime.py'
]

#################
#  Your classes #
#################
classes = ('nodule',)         # 若将来有多类，按 VOC <name> 全列出
num_classes = len(classes)

#####################
#  Your data paths  #
#####################
# 使用“零拷贝”自定义数据集类 ThyroidVOCDataset：train/test 列表中是相对图片路径
img_root   = '/data/lining/data/Structured_Dataset/Thyroid_Data/Comprehensive_data/picture/images'
xml_root   = '/data/caoxu/dataset/div-align-dg/voc_labels'
train_list = '/data/caoxu/dataset/div-align-dg/voc_sets/train.txt'
test_list  = '/data/caoxu/dataset/div-align-dg/voc_sets/test.txt'

################################
#  Keep model, tweak classes   #
################################
model = dict(
    roi_head=dict(
        bbox_head=dict(num_classes=num_classes)
    )
)

########################################
#  Dataset & pipelines for your data   #
########################################
custom_imports = dict(
    imports=['mmdet.datasets.thyroid_voc', 'mmdet.datasets.pipelines.oa_mix'],
    allow_failed_imports=False)
dataset_type = 'ThyroidVOCDataset'

img_norm_cfg = dict(mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)

train_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(type='LoadAnnotations', with_bbox=True),
    dict(type='Resize', img_scale=[(1333, 640), (1333, 800)], multiscale_mode='range', keep_ratio=True),
    dict(type='RandomFlip', flip_ratio=0.5),
    dict(type='Normalize', **img_norm_cfg),
    dict(type='Pad', size_divisor=32),
    dict(type='DefaultFormatBundle'),
    dict(type='Collect', keys=['img', 'gt_bboxes', 'gt_labels']),
]

test_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(
        type='MultiScaleFlipAug',
        img_scale=(1333, 800), flip=False,
        transforms=[
            dict(type='Resize', keep_ratio=True),
            dict(type='RandomFlip'),
            dict(type='Normalize', **img_norm_cfg),
            dict(type='Pad', size_divisor=32),
            dict(type='DefaultFormatBundle'),
            dict(type='Collect', keys=['img']),
        ])
]

data = dict(
    samples_per_gpu=2,     # 按显存调整
    workers_per_gpu=2,
    train=dict(
        type=dataset_type,
        ann_file=train_list,
        img_root=img_root,
        xml_root=xml_root,
        classes=classes,
        pipeline=train_pipeline),
    val=dict(
        type=dataset_type,
        ann_file=test_list,
        img_root=img_root,
        xml_root=xml_root,
        classes=classes,
        pipeline=test_pipeline),
    test=dict(
        type=dataset_type,
        ann_file=test_list,
        img_root=img_root,
        xml_root=xml_root,
        classes=classes,
        pipeline=test_pipeline),
)

################
#  Schedule    #
################
optimizer = dict(type='SGD', lr=0.02, momentum=0.9, weight_decay=0.0001)
optimizer_config = dict(grad_clip=None)

lr_config = dict(policy='step', step=[8, 11])
runner = dict(type='EpochBasedRunner', max_epochs=100)

evaluation = dict(interval=1, metric='mAP')
log_config = dict(interval=50, hooks=[dict(type='TextLoggerHook')])

